"""
@Time    :  2020/11/08 10:09
@Author  :  Sun_Z_Z  
@FileName:  RSA_Aug.py
@Institution:   Scut214
"""
import cv2
import numpy as  np
import numpy.random as random
from PIL import Image


# np.random.seed(2014)

class RandomShiftingAugmentation(object):
    """ RandomShiftingAugmentation with a given probability.
    You add this transform to Compose list directly.
    Args:
        p (float): probability of the image being RSA. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img):
        """
        Args:
            img (numpy Image): Image to be RSA.
        Returns:
            numpy Image: RandomShiftingAugmentation image.
        """
        if random.random() < self.p:
            return RSA_Augmentation(img)
        return img

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


def RSA_AugmentationA(img, r_c_min=0.7, r_h_min=0.7, r_w_min=0.7):
    H, W = img.shape[0:2]
    # random
    r_c = np.random.uniform(r_c_min, 1)
    r_start_y = np.random.uniform(0, 1 - r_c)
    r_h = np.random.uniform(r_h_min, 1 / r_c)
    r_w = np.random.uniform(r_w_min, 1)

    start_y, crop_H = int(r_start_y * H), int(r_c * H)
    img_crop = img[start_y:start_y + crop_H, :, :]
    re_H, re_W = int(r_h * img_crop.shape[0]), int(r_w * img_crop.shape[1])

    copy_start_y = np.random.randint(0, H - re_H)
    copy_start_x = np.random.randint(0, W - re_W)

    img_rescale = cv2.resize(img_crop, (re_W, re_H))
    img_pad = np.ones(img.shape, dtype=img.dtype) * 127
    np.copyto(img_pad[copy_start_y:copy_start_y + re_H, copy_start_x:copy_start_x + re_W], img_rescale)
    return img_pad


def RSA_AugmentationB(img, r_c_min=0.7, r_h_min=0.7, r_w_min=0.7):
    H, W = img.shape[0:2]
    r_h = np.random.uniform(r_h_min, 1)
    r_w = np.random.uniform(r_w_min, 1)
    re_H, re_W = int(r_h * H), int(r_w * W)
    img_rescale = cv2.resize(img, (re_W, re_H))

    copy_start_y = np.random.randint(0, H - re_H)
    copy_start_x = np.random.randint(0, W - re_W)

    img_pad = np.ones(img.shape, dtype=img.dtype) * 127
    np.copyto(img_pad[copy_start_y:copy_start_y + re_H, copy_start_x:copy_start_x + re_W], img_rescale)
    return img_pad


def RSA_Augmentation(img, p_crop=0.5, r_c_min=0.7, r_h_min=0.7, r_w_min=0.7):
    """
    p_crop: the probablity of  RSA_AugmentationA
    """
    img = np.array(img)
    if random.random() < p_crop:
        img_pad = RSA_AugmentationA(img, r_c_min, r_h_min, r_w_min)
    else:
        img_pad = RSA_AugmentationB(img, r_c_min, r_h_min, r_w_min)
    return Image.fromarray(img_pad)


if __name__ == '__main__':
    print('test RSA_Augmentation')
    img_path = './test1.jpg'
    img = cv2.imread(img_path)
    transform = RandomShiftingAugmentation()
    for i in range(20):
        img_pad = transform(img)
        print('input shape:', img.shape, 'output shape:', img_pad.shape)
        cv2.namedWindow('ret', 0)
        cv2.imshow('ret', img_pad)
        cv2.waitKey(0)
